########################################################
## Replication of Main Document
########################################################
rm(list=ls())
start_time <- Sys.time()
########################################################
## Package loading
########################################################
library(BridgeChange)
library(devtools)  
library(coda)
library(plm)
library(pforeach)
library(dplyr)
library(colorspace)
library(monomvn)
library(glmnet)
library(MASS)
library(MCMCpack)
library(pcse)
library(pastecs)
library(xtable)
library(bayesplot)
library(dotwhisker)
library(broom)
library(stableGR)
library(tidyverse)
library(lmtest)  #
source("helper.R")

########################################################################
## Application 1: Alvarez et al. data
########################################################################
data("agl")
model = "within"

set.seed(100033) 
formula <- growth ~ lagg1 + opengdp + openex + openimp + leftc + central
formula.inter <- growth ~ lagg1 + opengdp + openex + openimp + leftc + central + inter

mcmc = 1000; burn = 1000; verbose = 1000; thin = 1;
agl.cp0 <- BridgeFixedPanelPool(formula.p=formula, data.p = agl, mcmc.p=mcmc,
                                index.p = c('country', 'year'),
                                effect.p = 'time', standardize.p = TRUE,
                                Waic.p = TRUE, marginal.p = FALSE,
                                verbose.p=verbose, break.upper = 2,
                                interaction.p = 1)
agl.cp1 <- BridgeFixedPanelPool(formula.p=formula.inter, data.p = agl, mcmc.p=mcmc,
                                index.p = c('country', 'year'),
                                effect.p = 'time', standardize.p = TRUE,
                                Waic.p = TRUE, marginal.p = FALSE,
                                verbose.p=verbose, break.upper = 2,
                                interaction.p = 1)
agl.cp2 <- BridgeFixedPanelPool(formula.p=formula, data.p = agl, mcmc.p=mcmc,
                                index.p = c('country', 'year'),
                                effect.p = 'time', standardize.p = TRUE,
                                Waic.p = TRUE, marginal.p = FALSE,
                                verbose.p=verbose, break.upper = 2,
                                interaction.p = 2)

w0 <- attr(agl.cp0, "Waic")
w1 <- attr(agl.cp1, "Waic")
w2 <- attr(agl.cp2, "Waic")

#########################
## Table 1
#########################
tab1 <- xtable(matrix(unlist(rbind(w0, w1, w2)[,1]), 3, 3, byrow=TRUE), digit=0)

print(tab1, file="Table1.tex")
#########################
## Figure5
#########################
## Figure 5a: state transition
pdf(file="Figure5a.pdf", family="sans", height = 3, width = 12)
par(mfrow=c(1,4))
plotState(agl.cp1[[2]], legend.control =c(1970, 0.85), main="Partial interaction (break 1)", start = 1970)
plotState(agl.cp1[[3]], legend.control =c(1970, 0.85), main="Partial interaction (break 2)", start = 1970)
plotState(agl.cp2[[2]], legend.control =c(1970, 0.85), main="Full interaction (break 1)", start = 1970)
plotState(agl.cp2[[3]], legend.control =c(1970, 0.85), main="Full interaction (break 2)", start = 1970)
dev.off()

## basic model and three models
out <- three.plm(formula.inter, data=agl,
                 df.pre = subset(agl, year < 1979),
                 df.post = subset(agl, year >= 1979),
                 index = c('country', 'year'),
                 model = 'within', effect = 'time')
stargazer::stargazer(out[[1]], out[[2]], out[[3]], 
                     column.labels=c("Pooled FE","Regime 1 FE", "Regime 2 FE"),
                     align=TRUE)

## regime difference
output <- agl.cp1[[2]]
coefs <- attr(output, "hybrid") 
df.coef <- tibble(coef = rownames(coefs), Regime1 = coefs[,1], Regime2 = coefs[,2])
df.coef <- df.coef%>% mutate(diff = Regime2 - Regime1) %>% arrange(desc(abs(diff)))


## Figure 5b: shrinkage estimates
pdf(file="Figure5b.pdf", width=10, height = 5, family="sans")
par (mar=c(3,3,2,1), mgp=c(2,.7,0), tck=-.01)
par(mfrow=c(1, 2))
set.seed(99)
dotplotRegime(agl.cp1[[2]], hybrid=FALSE, start = 1970, text.cex=0.8, distance=6, pos = 3, x.location="default",
              main=expression(paste("Partial interaction model with one break")))
dotplotRegime(agl.cp1[[3]], hybrid=FALSE,  start = 1970, text.cex=0.8, distance=6, pos = 3, x.location="default",
              main=expression(paste("Partial interaction model with two breaks")))
dev.off()

#########################
## Figure 6
#########################
## 1. plm output
formula <-  formula.inter
agl.scale <- agl
agl.scale$growth <- scale(agl$growth)
agl.scale$lagg1 <- scale(agl$lagg1)
agl.scale$opengdp <- scale(agl$opengdp)
agl.scale$openex <- scale(agl$openex)
agl.scale$openimp <- scale(agl$openimp)
agl.scale$leftc <- scale(agl$leftc)
agl.scale$central <- scale(agl$central)
agl.scale$inter <- scale(agl$inter)

p.df.scaled <- pdata.frame(agl.scale, index = c('country', 'year'),
                 drop.index = FALSE) 
plm.fit <- plm(formula, data = p.df.scaled, model = 'within', effect = 'time')

data = agl
G <- length(unique(data$country))
N <- length(data$country)
dfa <- (G/(G - 1)) * (N - 1)/plm.fit$df.residual
## display with cluster VCE and df-adjustment
c_vcov <- dfa * vcovHC(plm.fit, type = "HC3", cluster = "group", adjust = TRUE)
coef.mat <- lmtest::coeftest(plm.fit, vcov = c_vcov)
## var.names <- rownames(coef.mat)
## rownames(coef.mat) <- gsub("X", "", var.names)

varnames <- stringr::str_replace(attr(output, "xnames")[,1], "_regime1","")

## 2. bayesian outputs
coef.out <- summary(output)[[1]][grep("beta", rownames(summary(output)[[1]])), ]
coef.hpd <- summary(output)[[2]][grep("beta", rownames(summary(output)[[2]])), ]

coef.out0 <- summary(agl.cp1[[1]])[[1]][grep("beta", rownames(summary(agl.cp1[[1]])[[1]])), ]
coef.hpd0 <- summary(agl.cp1[[1]])[[2]][grep("beta", rownames(summary(agl.cp1[[1]])[[2]])), ]


coef.regime0 <- as.data.frame(coef.out0)
coef.hpd0 <- as.data.frame(coef.hpd0)
rownames(coef.regime0) <- varnames


coef.regime1 <- as.data.frame(coef.out[grep("regime1", rownames(coef.out)), ])
coef.hpd1 <- as.data.frame(coef.hpd[grep("regime1", rownames(coef.hpd)), ])
rownames(coef.regime1) <- rownames(coef.mat) <- varnames

coef.regime2 <- as.data.frame(coef.out[grep("regime2", rownames(coef.out)), ])
coef.hpd2 <- as.data.frame(coef.hpd[grep("regime2", rownames(coef.hpd)), ])
rownames(coef.regime2) <- rownames(coef.mat) <- varnames

## 3. compare results
m0_df <- tibble(var = rownames(coef.regime0),
                estimate = coef.regime0$Mean,
                upper = coef.hpd0$`97.5%`,
                lower = coef.hpd0$`2.5%`)%>%
    mutate(model = "HMBB (no break)")
m1_df <- tibble(var = rownames(coef.mat),
                estimate = coef.mat[,1],
                upper = coef.mat[,1] + 1.96*coef.mat[,2],
                lower = coef.mat[,1] - 1.96*coef.mat[,2],) %>%
    mutate(model = "Fixed-effects")
m2_df <- tibble(var = rownames(coef.regime1),
                estimate = coef.regime1$Mean,
                upper = coef.hpd1$`97.5%`,
                lower = coef.hpd1$`2.5%`)%>%
    mutate(model = "HMBB Regime 1")
m3_df <- tibble(var = rownames(coef.regime2),
                estimate = coef.regime2$Mean,
                upper = coef.hpd2$`97.5%`,
                lower = coef.hpd2$`2.5%`)%>%
    mutate(model = "HMBB Regime 2")
three_models <- rbind(m1_df, m2_df, m3_df)

pdf(file="Figure6.pdf", width=7, height=5, family="sans")
ggplot(three_models, aes(x=var, y = estimate, col=model, shape=model)) +
    ## ggplot2 plot with confidence intervals
    geom_point(size = 2, position=position_dodge(width=0.5)) +
    geom_errorbar(
        aes(ymin = lower, ymax = upper),
        width = 0.1,
        position=position_dodge(width=0.5)) +
    geom_hline(aes(yintercept = 0), linetype = "dashed") +
    theme_bw()+
    theme(legend.position="bottom") + 
    coord_flip() + 
    labs(title="Fixed-effects estimates and HMBB estimates",
         x="Covariates", y="Estimate", col="Model", shape="Model")
dev.off()


#########################
## Table 3
#########################
## sparse comparison
out <- agl.cp2[[2]]
ns <- attr(agl.cp2[[2]], "m") + 1
X <- attr(agl.cp2[[2]], "X")
k <- ncol(X)
beta.target <- out[, grepl("beta", colnames(out))]
coef <- matrix(apply(beta.target, 2, mean), k, ns)
rownames(coef) <- colnames(X)    
sparse.coef <- attr(out, "hybrid")
tab3 <- xtable(cbind(coef[,1], sparse.coef[,1], coef[,2], sparse.coef[,2]))

print(tab3, file="Table3.tex")




end_time <- Sys.time()
print(end_time - start_time)
